#!/usr/bin/env python3
# military_target_detection.py

import cv2
import numpy as np
import onnxruntime
import threading
import time
import os
import queue

class MilitaryTargetDetection:
    def __init__(self, model_path, labels_path, play_device, logger=None, continuous_mode=False):
        # 加载标签
        self.labels = []
        with open(labels_path, 'r') as f:
            self.labels = [label.strip() for label in f.readlines()]
        
        # 设置推理会话选项
        session_options = onnxruntime.SessionOptions()
        session_options.intra_op_num_threads = 1
        
        # 创建ONNX Runtime会话
        self.session = onnxruntime.InferenceSession(
            model_path, 
            sess_options=session_options, 
            providers=["SpaceMITExecutionProvider"]
        )
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
        
        # 音频设备
        self.play_device = play_device
        
        # 日志记录器
        self.logger = logger
        
        # 持续运行模式
        self.continuous_mode = continuous_mode
        
        # 状态变量
        self.is_running = False
        self.detection_thread = None
        self.capture_thread = None
        self.cap = None
        self.frame_queue = queue.Queue(maxsize=1)
        self.result_queue = queue.Queue(maxsize=1)
        
        # 当前检测结果
        self.current_target = "无目标"
        self.current_confidence = 0.0
        
        print("军事目标检测模块初始化完成")
        
        # 如果设置为持续运行模式，自动启动
        if continuous_mode:
            self.start_detection()

    def preprocess(self, img):
        """预处理图像函数"""
        img = img / 255.
        img = cv2.resize(img, (256, 256))
        h, w = img.shape[0], img.shape[1]
        y0 = (h - 224) // 2
        x0 = (w - 224) // 2
        img = img[y0 : y0+224, x0 : x0+224, :]
        img = (img - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
        img = np.transpose(img, axes=[2, 0, 1])
        img = img.astype(np.float32)
        img = np.expand_dims(img, axis=0)
        return img

    def _detection_worker(self):
        """检测工作线程"""
        print("目标检测线程启动")
        while self.is_running:
            try:
                # 获取帧
                frame = self.frame_queue.get(timeout=0.5)
                
                # 预处理图像
                img = self.preprocess(frame)
                
                # 执行推理
                result = self.session.run([self.output_name], {self.input_name: img})[0]
                
                # 处理结果
                result = np.squeeze(result)
                top_k = result.argsort()[-5:][::-1]
                top_class = self.labels[top_k[0]]
                confidence = result[top_k[0]]
                
                # 更新当前结果
                self.current_target = top_class
                self.current_confidence = confidence
                
                # 将结果放入结果队列
                self.result_queue.put((top_class, confidence))
                
            except queue.Empty:
                continue
            except Exception as e:
                print(f"目标检测错误: {e}")
                continue
        
        print("目标检测线程退出")

    def find_camera_device(self):
        """查找可用的摄像头设备"""
        # 尝试直接打开端口20，使用V4L2后端
        try:
            cap = cv2.VideoCapture(20, cv2.CAP_V4L2)
            if cap.isOpened():
                print(f"成功打开端口20的摄像头")
                return cap
            cap.release()
        except:
            pass
        
        # 尝试打开默认摄像头，使用V4L2
        try:
            cap = cv2.VideoCapture(0, cv2.CAP_V4L2)
            if cap.isOpened():
                print(f"成功打开默认摄像头")
                return cap
            cap.release()
        except:
            pass
        
        # 尝试所有可能的设备
        for i in range(10):
            try:
                cap = cv2.VideoCapture(i, cv2.CAP_V4L2)
                if cap.isOpened():
                    print(f"成功打开摄像头索引: {i}")
                    return cap
                cap.release()
            except:
                continue
        
        # 尝试通过设备路径打开
        video_devices = [
            "/dev/video0", "/dev/video1", "/dev/video2", "/dev/video3",
            "/dev/video4", "/dev/video5", "/dev/video6", "/dev/video7",
            "/dev/video8", "/dev/video9", "/dev/video20"
        ]
        
        for device in video_devices:
            try:
                if os.path.exists(device):
                    cap = cv2.VideoCapture(device, cv2.CAP_V4L2)
                    if cap.isOpened():
                        print(f"成功打开摄像头设备: {device}")
                        return cap
                    cap.release()
            except:
                continue
        
        print("警告: 无法打开任何摄像头设备")
        return None

    def start_detection(self):
        """启动目标检测"""
        if self.is_running:
            print("目标检测已在运行中")
            return True
        
        # 查找并打开摄像头
        self.cap = self.find_camera_device()
        if self.cap is None:
            print("错误: 无法打开任何摄像头设备")
            return False
        
        # 设置摄像头参数
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        self.cap.set(cv2.CAP_PROP_FPS, 10)  # 降低帧率以提高稳定性
        
        # 启动检测线程
        self.is_running = True
        self.detection_thread = threading.Thread(
            target=self._detection_worker,
            daemon=True
        )
        self.detection_thread.start()
        
        # 启动摄像头采集线程
        self.capture_thread = threading.Thread(
            target=self._capture_worker,
            daemon=True
        )
        self.capture_thread.start()
        
        print("目标检测已启动")
        return True

    def _capture_worker(self):
        """摄像头采集线程"""
        print("摄像头采集线程启动")
        while self.is_running:
            try:
                # 读取摄像头帧
                ret, frame = self.cap.read()
                if not ret:
                    print("错误: 无法读取帧")
                    time.sleep(0.1)
                    continue
                
                # 将帧放入队列（如果队列不满）
                if self.frame_queue.empty():
                    try:
                        self.frame_queue.put(frame.copy(), block=False)
                    except queue.Full:
                        pass
                
            except Exception as e:
                print(f"摄像头采集错误: {e}")
                continue
        
        print("摄像头采集线程退出")

    def get_latest_result(self):
        """获取最新检测结果"""
        try:
            # 获取队列中的所有结果，只保留最新的
            while not self.result_queue.empty():
                self.current_target, self.current_confidence = self.result_queue.get(block=False)
        except queue.Empty:
            # 使用上一次的结果
            pass
        
        return self.current_target, self.current_confidence

    def speak_target(self):
        """语音播报当前检测到的目标"""
        target, confidence = self.get_latest_result()
        
        if "tank" in target.lower():
            wav_file = "tools/feedback_voice/faxian_tanke.wav"
        elif "armor" in target.lower():
            wav_file = "tools/feedback_voice/faxian_zhuangjia.wav"
        elif "soldier" in target.lower():
            wav_file = "tools/feedback_voice/faxian_shibing.wav"
        else:
            wav_file = "tools/feedback_voice/faxian_mubiao.wav"
        
        # 播放语音反馈
        if os.path.exists(wav_file):
            self.play_wav(wav_file)
        else:
            print(f"警告: 语音文件不存在 - {wav_file}")
        
        return target

    def play_wav(self, wav_file):
        """播放WAV文件"""
        # 在实际实现中，您应该使用您的音频播放功能
        # 这里使用简单的系统命令作为示例
        os.system(f"aplay -D {self.play_device} {wav_file}")

    def stop_detection(self):
        """停止目标检测"""
        if not self.is_running:
            return
        
        self.is_running = False
        
        # 等待线程结束
        if self.detection_thread and self.detection_thread.is_alive():
            self.detection_thread.join(timeout=1.0)
        
        if self.capture_thread and self.capture_thread.is_alive():
            self.capture_thread.join(timeout=1.0)
        
        # 释放摄像头
        if self.cap:
            self.cap.release()
            self.cap = None
        
        print("目标检测已停止")

    def __del__(self):
        try:
            # 确保安全地调用 stop_detection
            if hasattr(self, 'is_running') and self.is_running:
                self.stop_detection()
        except Exception as e:
            print(f"目标检测器销毁错误: {e}")